import pandas as pd
import numpy as np
import os

print("规则设定开始~")

# 如果当前路径没有"_2.月度组合"文件夹，则创建该文件夹
if not os.path.exists("_2.月度组合"):
    os.mkdir("_2.月度组合")

files = ["2024-12-31.xlsx", "2025-01-27.xlsx", "2025-02-28.xlsx", "2025-03-31.xlsx"]

for file in files:
    # --------------------------------------------------------------------------1.输入数据
    df = pd.read_excel("_1.截面数据\\" + file)

    # --------------------------------------------------------------------------2.处理数据
    # ---------------------------------------------------------------2.1选股问题
    """
    规则说明：
    1.月度涨跌幅：太高太低都不好，选取中间50%
    2.市盈率：通常认为低点好，选较小的50%
    3.净资产收益率ROE：通常认为高点好，选较大的50%
    4.这个策略是做着玩的，所以就从剩下的股票里随机选30个，避免每个人跑出来一样的结果
    """

    req1_l = df["月度涨跌幅"].quantile(25 / 100)
    req1_u = df["月度涨跌幅"].quantile(75 / 100)
    req2 = df["市盈率"].quantile(50 / 100)
    req3 = df["净资产收益率ROE"].quantile(50 / 100)

    df = df.drop(index=(df.loc[(df["月度涨跌幅"] < req1_l)].index))
    df = df.drop(index=(df.loc[(df["月度涨跌幅"] > req1_u)].index))
    df = df.drop(index=(df.loc[(df["市盈率"] > req2)].index))
    df = df.drop(index=(df.loc[(df["净资产收益率ROE"] < req3)].index))

    rand_num = np.random.rand(len(df)) * 100  # 生成一列0~100之间的随机数
    rand_num = np.round(rand_num, decimals=4)  # 保留4位小数
    df["随机序号"] = rand_num  # 随机序号列
    df = df.sort_values(by="随机序号", ascending=False)  # 降序
    df.index = range(len(df))  # 重置索引
    df = df.drop(index=(df.loc[(df.index >= 10)].index))  # 选取前10个股票作为的投资组合

    # ---------------------------------------------------------------2.2择时问题
    file_year = int(file.split("-")[0])
    file_month = int(file.split("-")[1])

    # 买入时间：第二天（下个月第一天）
    # 卖出时间：一个月后（下下个月第一天）
    if file_month == 12:
        date_buy = pd.Timestamp(file_year + 1, 1, 1).date()
        date_sell = pd.Timestamp(file_year + 1, 2, 1).date()
    else:
        date_buy = pd.Timestamp(file_year, file_month + 1, 1).date()
        date_sell = pd.Timestamp(file_year, file_month + 2, 1).date()

    df["买入时间"] = date_buy
    df["卖出时间"] = date_sell

    # ---------------------------------------------------------------2.3权重问题
    df["金额比例"] = 1 / 10  # 等权重

    # --------------------------------------------------------------------------3.输出策略
    df = df.drop(columns=["月度涨跌幅", "市盈率", "净资产收益率ROE", "随机序号"])  # 删除多余列
    df.insert(loc=0, column="序号", value=range(1, len(df) + 1))  # 在第一列加一列序号

    df.to_excel("_2.月度组合\\" + str(date_buy) + "(组合).xlsx", index=False)

print("规则设定完成！")
